import os
import torch
import numpy as np
from PIL import Image, ImageOps, ImageSequence
import folder_paths
from comfy.model_management import get_torch_device
import comfy.utils
import datetime

def create_default_image_tensor(width=512, height=512, reference_images=None):
    """
    创建默认图像张量，支持动态尺寸
    
    Args:
        width: 指定宽度
        height: 指定高度
        reference_images: 参考图像列表，用于确定尺寸
    
    Returns:
        torch.Tensor: 默认图像张量 (1, H, W, 3) float32 [0,1]
    """
    # 如果有参考图像，使用第一个图像的尺寸
    if reference_images and len(reference_images) > 0:
        ref_img = reference_images[0]
        if isinstance(ref_img, torch.Tensor) and len(ref_img.shape) >= 3:
            height = ref_img.shape[-3] if len(ref_img.shape) == 4 else ref_img.shape[-3]
            width = ref_img.shape[-2] if len(ref_img.shape) == 4 else ref_img.shape[-2]
    
    # 创建纯黑色图像张量
    image_tensor = torch.zeros((1, height, width, 3), dtype=torch.float32)
    return image_tensor

class ImageInfoNode:
    """
    图片信息处理器节点
    处理文件夹中的图像，提取信息并支持批次合并
    """
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "folder_path": ("STRING", {"default": "", "multiline": False, "tooltip": "图像文件夹路径"}),
            },
            "optional": {
                "batch_resize_mode": (["first_image", "max_size", "min_size"], 
                                    {"default": "first_image", "tooltip": "批次图像尺寸统一模式"}),
            },
        }

    RETURN_TYPES = ("LIST", "LIST", "LIST", "LIST", "LIST", "IMAGE", "IMAGE", "INT")
    RETURN_NAMES = ("images", "image_paths", "image_names", "names_without_ext", 
                    "image_structures", "preview", "batch", "total_count")
    FUNCTION = "process_images"
    CATEGORY = "AI训练/image"

    def process_images(self, folder_path, batch_resize_mode="first_image"):
        """
        处理图像文件夹
        """
        try:
            if not folder_path or not os.path.exists(folder_path):
                default_img = create_default_image_tensor()
                return ([], [], [], [], [], default_img, default_img, 0)

            # 支持的图像格式
            supported_formats = {'.png', '.jpg', '.jpeg', '.bmp', '.tiff', '.tga', '.webp'}
            
            images = []
            image_paths = []
            image_names = []
            names_without_ext = []
            image_structures = []

            # 遍历文件夹
            for filename in os.listdir(folder_path):
                file_path = os.path.join(folder_path, filename)
                
                if os.path.isfile(file_path):
                    file_ext = os.path.splitext(filename)[1].lower()
                    
                    if file_ext in supported_formats:
                        try:
                            # 加载图像
                            pil_image = Image.open(file_path)
                            
                            # 处理EXIF方向
                            pil_image = ImageOps.exif_transpose(pil_image)
                            
                            # 处理多帧图像
                            if hasattr(pil_image, 'n_frames') and pil_image.n_frames > 1:
                                pil_image = pil_image.copy()
                            
                            # 转换为RGB
                            if pil_image.mode != 'RGB':
                                pil_image = pil_image.convert('RGB')
                            
                            # 转换为张量
                            image_np = np.array(pil_image).astype(np.float32) / 255.0
                            image_tensor = torch.from_numpy(image_np).unsqueeze(0)  # (1, H, W, 3)
                            
                            # 收集信息
                            name_without_ext = os.path.splitext(filename)[0]
                            folder_path_abs = os.path.dirname(file_path)
                            
                            # 创建图像结构
                            image_structure = {
                                "filename": filename,
                                "name_without_ext": name_without_ext,
                                "file_path": file_path,
                                "folder_path": folder_path_abs,
                                "width": pil_image.width,
                                "height": pil_image.height,
                                "image_tensor": image_tensor
                            }
                            
                            images.append(image_tensor)
                            image_paths.append(file_path)
                            image_names.append(filename)
                            names_without_ext.append(name_without_ext)
                            image_structures.append(image_structure)

                        except Exception as e:
                            continue

            if not images:
                default_img = create_default_image_tensor()
                return ([], [], [], [], [], default_img, default_img, 0)

            # 创建预览图像（第一张）
            preview = images[0]
            
            # 创建批次图像
            batch = self._create_batch_images(images, batch_resize_mode)
            
            total_count = len(images)

            return (images, image_paths, image_names, names_without_ext,
                   image_structures, preview, batch, total_count)
            
        except Exception as e:
            default_img = create_default_image_tensor()
            return ([], [], [], [], [], default_img, default_img, 0)

    def _create_batch_images(self, images, resize_mode):
        """创建批次图像"""
        if not images:
            return create_default_image_tensor()
        
        if len(images) == 1:
            return images[0]
        
        # 确定目标尺寸
        if resize_mode == "first_image":
            target_height = images[0].shape[1]
            target_width = images[0].shape[2]
        elif resize_mode == "max_size":
            target_height = max(img.shape[1] for img in images)
            target_width = max(img.shape[2] for img in images)
        elif resize_mode == "min_size":
            target_height = min(img.shape[1] for img in images)
            target_width = min(img.shape[2] for img in images)
        else:
            target_height = images[0].shape[1]
            target_width = images[0].shape[2]
        
        # 调整所有图像到目标尺寸
        resized_images = []
        for img in images:
            if img.shape[1] != target_height or img.shape[2] != target_width:
                # 使用ComfyUI的缩放函数
                resized = comfy.utils.common_upscale(
                    img.movedim(-1, 1),  # (1, 3, H, W)
                    target_width, target_height, 
                    "lanczos", "center"
                ).movedim(1, -1)  # (1, H, W, 3)
                resized_images.append(resized)
            else:
                resized_images.append(img)
        
        # 合并为批次
        batch = torch.cat(resized_images, dim=0)

        return batch

class BatchImageNode:
    """
    图片批次处理器节点
    将图像列表合并为批次张量
    """
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "images": ("LIST", {"tooltip": "图像列表"}),
            },
        }

    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "batch_images"
    CATEGORY = "AI训练/image"

    def batch_images(self, images):
        if not images:
            return (create_default_image_tensor(),)
        
        if len(images) == 1:
            return (images[0],)
        
        # 合并所有图像为批次
        batch = torch.cat(images, dim=0)
        return (batch,)

class ImageSelectorNode:
    """
    图片选择器节点
    从图像列表中选择指定索引的图像
    """
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "images": ("LIST", {"tooltip": "图像列表"}),
                "index": ("INT", {"default": 0, "min": 0, "max": 9999, "tooltip": "要选择的图像索引"}),
            },
        }

    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "select_image"
    CATEGORY = "AI训练/image"

    def select_image(self, images, index):
        if not images or index >= len(images):
            return (create_default_image_tensor(reference_images=images),)
        
        return (images[index],)

class LoadImageFromPath:
    """
    从路径加载图像节点
    根据指定的图像路径加载图像
    """
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image_path": ("STRING", {"default": "", "multiline": False, "tooltip": "图像文件路径"}),
            },
        }

    RETURN_TYPES = ("IMAGE", "MASK")
    RETURN_NAMES = ("image", "mask")
    FUNCTION = "load_image"
    CATEGORY = "AI训练/image"

    def load_image(self, image_path):
        try:
            # 获取完整路径
            if os.path.exists(image_path):
                full_path = image_path
            else:
                # 尝试在输入目录中查找
                input_dir = folder_paths.get_input_directory()
                full_path = os.path.join(input_dir, image_path)
                
                if not os.path.exists(full_path):
                    default_img = create_default_image_tensor()
                    default_mask = torch.zeros((1, default_img.shape[1], default_img.shape[2]), dtype=torch.float32)
                    return (default_img, default_mask)
            
            # 加载图像
            pil_image = Image.open(full_path)
            
            # 处理EXIF方向
            pil_image = ImageOps.exif_transpose(pil_image)
            
            # 处理透明度
            output_images = []
            output_masks = []
            
            # 处理多帧图像
            for frame in ImageSequence.Iterator(pil_image):
                frame = frame.copy()
                
                # 创建遮罩
                if 'transparency' in frame.info:
                    mask = np.array(frame.convert('RGBA'))[:,:,3].astype(np.float32) / 255.0
                elif frame.mode == 'RGBA':
                    mask = np.array(frame)[:,:,3].astype(np.float32) / 255.0
                else:
                    mask = np.ones((frame.height, frame.width), dtype=np.float32)
                
                # 转换为RGB
                if frame.mode != 'RGB':
                    frame = frame.convert('RGB')
                
                # 转换为张量
                image_np = np.array(frame).astype(np.float32) / 255.0
                image_tensor = torch.from_numpy(image_np).unsqueeze(0)  # (1, H, W, 3)
                mask_tensor = torch.from_numpy(mask).unsqueeze(0)  # (1, H, W)
                
                output_images.append(image_tensor)
                output_masks.append(mask_tensor)
            
            if len(output_images) > 1:
                # 多帧图像，合并为批次
                image_batch = torch.cat(output_images, dim=0)
                mask_batch = torch.cat(output_masks, dim=0)
            else:
                image_batch = output_images[0]
                mask_batch = output_masks[0]
            
            return (image_batch, mask_batch)

        except Exception as e:
            default_img = create_default_image_tensor()
            default_mask = torch.zeros((1, default_img.shape[1], default_img.shape[2]), dtype=torch.float32)
            return (default_img, default_mask)

class ImageStructureSelector:
    """
    图像结构选择器节点
    从image_structures列表中选择指定索引的结构，输出其所有值
    """
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image_structures": ("LIST", {"tooltip": "图像结构列表"}),
                "index": ("INT", {"default": 0, "min": 0, "max": 9999, "tooltip": "要选择的结构索引"}),
            },
        }

    RETURN_TYPES = ("STRING", "STRING", "STRING", "STRING", "INT", "INT", "IMAGE")
    RETURN_NAMES = ("filename", "name_without_ext", "file_path", "folder_path", "width", "height", "image_tensor")
    FUNCTION = "select_structure"
    CATEGORY = "AI训练/image"

    def select_structure(self, image_structures, index):
        """
        选择指定索引的图像结构并输出其所有值
        
        Args:
            image_structures: 图像结构列表
            index: 要选择的索引
            
        Returns:
            tuple: (filename, name_without_ext, file_path, folder_path, width, height, image_tensor)
        """
        try:
            # 运行时验证
            if not isinstance(image_structures, list):
                return self._get_default_structure()

            # 验证index类型
            if not isinstance(index, int):
                try:
                    index = int(index)
                except (ValueError, TypeError):
                    index = 0

            # 验证输入
            if not image_structures:
                return self._get_default_structure()

            # 处理索引越界
            if index < 0:
                index = 0
            elif index >= len(image_structures):
                index = len(image_structures) - 1
            
            # 获取选中的结构
            selected_structure = image_structures[index]
            
            # 提取所有值
            filename = selected_structure.get("filename", "unknown.png")
            name_without_ext = selected_structure.get("name_without_ext", "unknown")
            file_path = selected_structure.get("file_path", "")
            folder_path = selected_structure.get("folder_path", "")
            width = selected_structure.get("width", 512)
            height = selected_structure.get("height", 512)
            image_tensor = selected_structure.get("image_tensor", create_default_image_tensor(width, height))
            
            return (filename, name_without_ext, file_path, folder_path, width, height, image_tensor)

        except Exception as e:
            return self._get_default_structure()
    
    def _get_default_structure(self):
        """获取默认结构"""
        default_tensor = create_default_image_tensor()
        return ("default.png", "default", "", "", 512, 512, default_tensor)

# 节点映射
NODE_CLASS_MAPPINGS = {
    "BatchImageNode": BatchImageNode,
    "ImageSelectorNode": ImageSelectorNode,
    "LoadImageFromPath": LoadImageFromPath,
    "ImageStructureSelector": ImageStructureSelector
}

# 显示名称映射
NODE_DISPLAY_NAME_MAPPINGS = {
    "BatchImageNode": "图像批次处理器",
    "ImageSelectorNode": "图像选择器",
    "LoadImageFromPath": "从路径加载图像",
    "ImageStructureSelector": "图像结构选择器"
}
